import json
import os
import time

import chromadb
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from chromadb.utils import embedding_functions

from app.provider.model_provider import ModelProvider
from zhanshop.app import App
from zhanshop.env import Env

class ImgsearchProvider(ModelProvider):
    modelParam = None
    chromaClient = None
    tableName = None
    model = None
    collection = None
    transform = None
    def __init__(self, dbName, tableName, modelParam: json):
        self.dbName = dbName
        self.tableName = tableName
        self.modelParam = modelParam
        dbPath = App.make(Env).rootPath + "/runtime/chroma/" + dbName
        self.chromaClient = chromadb.PersistentClient(path=dbPath)

        """模型初始化"""
        model = models.resnet50(weights=None)
        model.load_state_dict(torch.load(self.getResNet50ModelPath(models.ResNet50_Weights.IMAGENET1K_V2)))
        model = torch.nn.Sequential(*(list(model.children())[:-1]))
        model.eval()
        self.model = model

        # model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        # model = torch.nn.Sequential(*(list(model.children())[:-1]))  # 移除最后一层分类层
        # model.eval()  # 设置为评估模式
        # self.model = model

        """图像预处理管道"""
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def getCollection(self):
        """
        获取图像合集
        :return:
        """
        if (self.collection == None):
            sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
                model_name=self.getSentenceTransformerModelPath(self.modelParam["model"]),
            )

            self.modelParam["chromadb"]["name"] = self.tableName
            self.modelParam["chromadb"]['embedding_function'] = sentence_transformer_ef
            chromadbParam = self.modelParam["chromadb"]
            collection = self.chromaClient.get_or_create_collection(**chromadbParam)
            self.collection = collection

        return self.collection

    def vectorCode(self, imagePath):
        """
        获取图片向量
        :param imagePath:
        :param transform:
        :return:
        """
        img = Image.open(imagePath).convert('RGB')
        model = self.model
        transform = self.transform
        img_t = transform(img).unsqueeze(0)  # 增加batch维度

        with torch.no_grad():
            features = model(img_t)

        # 展平并转换为numpy数组
        return features.squeeze().numpy().tolist()

    def batchAdd(self, _imagePaths: list, _metadatas: list = None, _uris: list = None):
        """
        批量添加数据
        :param ids: 向量ID列表
        :param documents: 原始文档列表
        :param metadatas: 关联的元数据
        :param uris: 文档内容的URI
        :return:
        """
        vectorCodes = []
        _ids = []
        prefix = time.time_ns() // 1000
        for key,imagePath in enumerate(_imagePaths):
            vectorCode = self.vectorCode(imagePath)
            vectorCodes.append(vectorCode)
            _ids.append(str(prefix + key))

        # 将向量数据和原始文档数据写入数据库
        # 如果有值就不要写入了
        self.getCollection().add(
            documents=_imagePaths,
            embeddings=vectorCodes,
            ids=_ids,
            metadatas=_metadatas,
            uris=_uris
        )

    def query(self, imgPath: str, limit: int = 10000, isOriginal = False):
        """
        查询相似图片
        :param imgPath: 检索的图片
        :param limit: 查询条数
        :param isOriginal: 是否返回原始内容
        :return:
        """
        # 提取查询图像特征
        vectorCode = self.vectorCode(imgPath)

        # 执行查询
        results = self.getCollection().query(
            query_embeddings=[vectorCode],
            n_results=limit
        )

        datas = {"list": [], "included": results["included"]}
        for key,document in enumerate(results["documents"][0]):
            datas["list"].append({
                "id": results["ids"][0][key],
                "document": document,
                "metadata": results["metadatas"][0][key],
                "distance": results["distances"][0][key],
            })

        return datas